# -*- coding: utf-8 -*-
# @Time    : 2019/5/13 14:36
# @Author  : zhouqiang
# @Site    :
# @File    : checker.py
# @Software: PyCharm

from quant_researcher.quant.project_tool.logger.my_logger import LOG
from quant_researcher.quant.project_tool import hammer

import pandas as pd


def check_length(df, inplace=False):
    len_before = df.shape[0]
    if inplace:
        df.dropna(inplace=True)  # 这里应该是值引用，会直接修改源数据
        len_after = df.shape[0]
    else:
        tmp = df.dropna()
        len_after = tmp.shape[0]
    if len_after != len_before:
        LOG.error(f'数据drop na前个数={len_before}，后个数={len_after}')
    else:
        LOG.info(f'数据drop na前个数=后个数={len_after}')


def null_or_na_checker(df, check_null=True, check_na=True, **kwargs):
    """
    判断 pd.DataFrame 中是不是有 null 或者 na，有则打印并报错

    :raise RuntimeError: 如果检查 null 或 na 且数据存在该种情况
    :param pd.DataFrame df: 被检查的数据
    :param bool check_null: 是否检查 null，默认：True
    :param bool check_na: 是否检查 na，默认：True
    :param kwargs:
        - inplace，bool，是否 dropna，默认：False
        - raise_err，bool，是否 raise_err，默认：True
    :return:
    """
    inplace = kwargs.pop('inplace', False)
    raise_err = kwargs.pop('raise_err', False)

    if check_null and pd.isnull(df).sum().sum() != 0:
        hammer.slim('数据是 null 的地方：')
        # TODO 这如果一行有多个空值，会重复打印这一行，可在 isnull 前加 [指定列]
        hammer.slim(df[df.isnull().values], a_f=1)
        if inplace or not raise_err:
            pass
        else:
            raise RuntimeError('数据存在 null')
    else:
        LOG.info('没找到 null')

    if check_na and df.isna().sum().sum() != 0:
        hammer.slim('数据是 na 的地方：')
        hammer.slim(df[pd.isna(df).values], a_f=1)
        if inplace or not raise_err:
            pass
        else:
            raise RuntimeError('数据存在 na')
    else:
        LOG.info('没找到 na')

    if inplace:
        df.dropna(inplace=True)  # 这里应该是值引用，会直接修改源数据
